from kuai.sim import setup_model, MolecularDynamicJob
from kuai.mol.io import read_molfile
from kuai.sim.dff2kuai import read_ppffile
from kuai.sim.efunc import make_efunc
from kuai.unit import default_units
from numpy import dot, sum
from math import sqrt
from kuai.kuaisuan import PBC

index, parameters = read_ppffile('C2H5OH and Acetone.ppf')

mol1 = read_molfile('C2H5OH.msd')
mol2 = read_molfile('Acetone.msd')

model = setup_model([(mol1, 100), (mol2, 100)], index, parameters)

coords = []

A, _, _ = default_units.parse(1, 'A')
print A

with open('C2H5OH and Acetone.msd') as file:
    for i in file:
        if i.startswith('PBC:'):
            model.pbc = PBC.parse(i.strip())
            model.pbc *= A
        else:
            try:
                natoms = int(i.strip())
                break
            except ValueError:
                continue
    assert natoms == len(mol1.atoms) * 100 + len(mol2.atoms) * 100
    for i in range(natoms):
        tokens = file.next().split()
        assert len(tokens) == 12
        coords.append(float(tokens[6])*A)
        coords.append(float(tokens[7])*A)
        coords.append(float(tokens[8])*A)
model.coords = coords

funcs = make_efunc(model, index)

for i in funcs:
    if i.name == 'pair-nonbond/lj12_6':
        i.cutoff = 10 * A       # Cutoff = 10 A

fe, _, _ = default_units.format(1, 'kcal/mol')
fg, _, _ = default_units.format(1, 'N')

job = MolecularDynamicJob()
e0 = job.getE(model, funcs)
e1, g1 = job.getEG(model, funcs)

for i in funcs:
    print i.name
print e0 * fe, sum(e0) * fe
print e1 * fe, sum(e1) * fe
print g1 * fg


for i in range(1, 100):
    job.setCPU(i)
    e2, g2 = job.getEG(model, funcs)
    deltae = e2-e1
    deltae = sqrt(dot(deltae, deltae))
    deltag = g2-g1
    deltag = sqrt(dot(deltag, deltag))
    print i, deltae*fe, deltae/sum(e1) * 100, deltag*fg

"""
model.potential_e = e1
model.kinetic_e = sum(e1)


from kuai.sim import SimpleReporter
reporter = SimpleReporter()
reporter.start(model)
for i in range(0, 100000, 1000):
    reporter.act(model, i+1)
reporter.finish(model)


"""
